""""""
import time
import cv2
import numpy as np
from ultralytics import YOLO

# ✅ 1. 视频文件路径配置
VIDEO_INPUT = r"F:\协助\协助-杨\250325红外路口人脸识别\20250325\20250325\人脸识别\1816437285179953152181735212_20250324171953-00.08.52.327-00.19.15.821-seg2-00.07.49.369-00.10.24.283-seg3.mp4"  # 输入视频文件路径
VIDEO_OUTPUT = r"F:\协助\协助-杨\250325红外路口人脸识别\识别结果\人脸\人脸识别结果.mp4"  # 输出视频文件路径

# ✅ 2. 加载 YOLO 模型
model = YOLO("D:/code_work/git/police_ai/weight/yolo_face.pt").to("cuda")

# ✅ 3. 读取视频文件
cap = cv2.VideoCapture(VIDEO_INPUT)
if not cap.isOpened():
    print("❌ 无法打开视频文件")
    exit()

# ✅ 4. 获取视频属性
width = int(cap.get(cv2.CAP_PROP_FRAME_WIDTH))
height = int(cap.get(cv2.CAP_PROP_FRAME_HEIGHT))
fps = int(cap.get(cv2.CAP_PROP_FPS))
total_frames = int(cap.get(cv2.CAP_PROP_FRAME_COUNT))

# ✅ 5. 设置视频写入器
fourcc = cv2.VideoWriter_fourcc(*'mp4v')
out = cv2.VideoWriter(VIDEO_OUTPUT, fourcc, fps, (width, height))

# ✅ 6. 处理视频帧
frame_count = 0
while cap.isOpened():
    ret, frame = cap.read()
    if not ret:
        print("❌ 视频读取完成")
        break

    # 显示处理进度
    frame_count += 1
    progress = (frame_count / total_frames) * 100
    print(f"\r处理进度: {progress:.2f}%", end="")

    # ✅ 目标检测
    results = model.predict(frame, imgsz=(width, height))
    
    # ✅ 获取检测框数量
    num_detections = len(results[0].boxes)
    
    # ✅ 绘制识别框
    annotated_frame = results[0].plot()
    
    # ✅ 在右上角显示检测框数量
    text = f"count: {num_detections}"
    # 获取文本大小以计算位置
    (text_width, text_height), _ = cv2.getTextSize(text, cv2.FONT_HERSHEY_SIMPLEX, 0.7, 2)
    # 计算右上角坐标，留出一些边距
    x = width - text_width - 10
    y = text_height + 10
    cv2.putText(annotated_frame, text, 
                (x, y), cv2.FONT_HERSHEY_SIMPLEX, 0.7, (0, 255, 0), 2)

    # ✅ 写入处理后的帧
    out.write(annotated_frame)

print("\n✅ 视频处理完成！")

# ✅ 7. 释放资源
cap.release()
out.release()
